import sys, os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
import json
import csv
from preprossData.utils.crop_padding import crop_resize_save
from tqdm import tqdm

def parse_csv(csv_path):
    """
    读取 CSV 文件，并返回一个字典，格式如下：
    
    {
       image_id: "<diagnosis text>"
    }
    
    其中，文本描述的生成规则为：
      - 对于 adjudicated_dr_grade（或 level，ICDR 5 级）:
           0 -> normal
           1 -> Mild Diabetic Retinopathy
           2 -> Moderate Diabetic Retinopathy
           3 -> Severe Diabetic Retinopathy
           4 -> Proliferative Diabetic Retinopathy
           
      - 如果 CSV 文件中既没有 "level" 也没有 "DR_grade" 列，则直接抛出异常。
      
    示例 CSV 文件内容：
        ,image,quality,DR_grade
        0,1_right.jpeg,1,0
    """
    dr_mapping = {
        "0": "normal",
        "1": "Mild Diabetic Retinopathy",
        "2": "Moderate Diabetic Retinopathy",
        "3": "Severe Diabetic Retinopathy",
        "4": "Proliferative Diabetic Retinopathy"
    }
    result = {}
    with open(csv_path, newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        fieldnames = reader.fieldnames or []
        # 检查使用哪一列来表示 DR 等级
        if "level" in fieldnames:
            grade_key = "level"
        elif "DR_grade" in fieldnames:
            grade_key = "DR_grade"
        else:
            raise ValueError("CSV file must contain either 'level' or 'DR_grade' column.")
        
        for row in reader:
            image_id = row.get("image", "").strip()
            dr_val = row.get(grade_key, "").strip()
            if not image_id or not dr_val:
                continue
            result[image_id] = dr_mapping.get(dr_val, dr_val)
    return result


def merge_data_dict(train_official, test_custom):
    """
    合并两个由 CSV 解析后的数据字典：
      - 如果 test_custom 中的某个 image_name 在 train_official 中不存在，则将该键值对加入 train_official。
    返回合并后的字典。
    """
    merged = train_official.copy()
    for k, v in test_custom.items():
        if k not in merged:
            merged[k] = v
    return merged

def gather_data(data_path, tar_path, prefix='eyepacs_', resize=(224,224)):
    """
    对新数据集进行预处理：
      - 数据集根目录下要求包含：
            trainLabels/trainLabels.csv
            Label_EyeQ_test.csv
            以及存放图像的 train 和 test 文件夹。
      - 分别解析 trainLabels 和测试 CSV，生成的标签字典 key 为图像文件名（含扩展名），值为诊断描述，
        然后将二者合并。
      - 遍历 train 和 test 文件夹，对每张图像调用 crop_resize_save 进行裁剪与 resize，
        并生成新的图像名称，格式为 "{prefix}_{原始文件名}.png"（保留原始文件名，仅修改后缀并加前缀）。
      - 同时生成一个 split 字典，记录各 split 下的处理后图像名称。
      - 最终生成的标注格式为：
            {
                new_image_name: {
                    "image_name": new_image_name,
                    "image_path": "images/{new_image_name}",
                    "diagnosis": {
                        "text": <diagnosis text>
                    }
                },
                ...
            }
      - 将标注信息保存到 tar_path/annotations.json，
        将 split 信息保存到 tar_path/split/official.json。
    
    参数：
        data_path (str): 新数据集根目录，要求包含 IMAGES 文件夹和相应 CSV 文件。
        tar_path (str): 预处理后数据存放目录。
        prefix (str): 处理后图片名称的前缀，默认为 "eyepacs_"
        resize (tuple): 目标尺寸，默认为 (224,224)
    
    返回：
        dict: 标注信息字典
    """
    # 创建目标目录
    os.makedirs(tar_path, exist_ok=True)
    images_out_dir = os.path.join(tar_path, "images")
    os.makedirs(images_out_dir, exist_ok=True)
    split_dir = os.path.join(tar_path, "split")
    os.makedirs(split_dir, exist_ok=True)
    
    # 解析 CSV 文件生成标签字典
    train_csv_path = os.path.join(data_path, "trainLabels", "trainLabels.csv")
    test_csv_path = os.path.join(data_path, "Label_EyeQ_test.csv")
    train_official = parse_csv(train_csv_path)
    test_custom = parse_csv(test_csv_path)
    label_dict = merge_data_dict(train_official, test_custom)
    
    final_annotations = {}
    official_split = {"train": [], "test": []}
    split_list = ["train", "test"]
    
    # 遍历每个 split 文件夹
    for split in split_list:
        image_src_dir = os.path.join(data_path, split)
        image_list=sorted(os.listdir(image_src_dir))
        total_image_num=len(image_list)
        for image_file in tqdm(image_list,desc=f'{split} process',unit='images',total=total_image_num):
            if not image_file.lower().endswith(('.png','.jpg','.jpeg','.bmp','.ppm')):
                continue
            # 保留原始文件名（不含扩展名），仅加前缀并统一保存为 PNG 格式
            base_name = os.path.splitext(image_file)[0]
            new_image_name = f"{prefix}_{base_name}.png"
            dest_image_path = os.path.join(images_out_dir, new_image_name)
            # 进行图像预处理：裁剪、resize 并保存为 PNG
            crop_info=crop_resize_save(
                image_path = os.path.join(image_src_dir, image_file),
                save_path = dest_image_path,
                resize = resize,
                crop_threshold = 15
            )
            # 获取对应的诊断描述（以原始文件名作为 key）
            diag_text = label_dict.get(image_file, "")
            final_annotations[new_image_name] = {
                "image_name": new_image_name,
                "image_path": os.path.join("images", new_image_name),
                'original_path':os.path.relpath(os.path.join( os.path.join(image_src_dir, image_file)),data_path),
                'crop_info':crop_info,
                "diagnosis": {
                    "text": diag_text
                }
            }
            official_split[split].append(new_image_name)
    
    # 保存标注信息
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(final_annotations, f, indent=4)
    
    # 保存分割信息
    official_split_path = os.path.join(split_dir, "official.json")
    with open(official_split_path, "w", encoding="utf-8") as f:
        json.dump(official_split, f, indent=4)
    
    return final_annotations

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Eyepacs 数据集预处理")
    parser.add_argument("--data_path", type=str, default="../Dataset/eyepacs",
                        help="新数据集根目录，要求其中包含 trainLabels, Label_EyeQ_test CSV 文件以及 train 和 test 图像文件夹")
    parser.add_argument("--tar_path", type=str, default="../Dataset/processed224_eyepacs",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="eyepacs_",
                        help="处理后图片名称的前缀，默认为 'eyepacs_'")
    parser.add_argument("--resize", type=int, nargs=2, default=[224,224],
                        help="目标尺寸，默认为 224 224")
    args = parser.parse_args()
    
    annotations = gather_data(args.data_path, args.tar_path, prefix=args.prefix, resize=tuple(args.resize))
    print("Preprocessing completed. Annotations saved.")
